Skip to content

Finalize pr2 fixes review#3

Closed
sheepdestroyer wants to merge 21 commits into
masterfrom
finalize-pr2-fixes-review
Closed

Finalize pr2 fixes review#3
sheepdestroyer wants to merge 21 commits into
masterfrom
finalize-pr2-fixes-review

Conversation

@sheepdestroyer

@sheepdestroyer sheepdestroyer commented Jun 18, 2026

Copy link
Copy Markdown
Owner

Summary of Completed Work
Dataset Visualizer Fixes (visualizer.html):

Fixed the Agreement Label Bug by updating .textContent on the inner .val span rather than replacing the entire container, preserving the UI labels.
Fixed the UI State Desync by calling render() inside saveNote(), giving immediate visual feedback (e.g., showing the REVIEWED tag) on note saves.
Cleaned up the dead code variable effectiveTier.
Extensible Breaker Mapping (agy_proxy.py):

Replaced the hardcoded index check (actual_tier_idx == 0) with dynamic model family matching ("gemini" in model_name), preventing future reordering failures.
Container Build Restoration (Containerfile):

Removed the gitignored best_free_model.json from the build's COPY sequence, resolving clean-checkout build failures.
100% Docstring Coverage:

Documented all missing module headers, class initializations, generators, and callback functions across main.py, agy_proxy.py, circuit_breaker.py, and the script utilities. This ensures the CI/CD docstring validation check will pass.
Automated & Manual Verification:

Ran all tests (test_circuit_breaker.py, verify_breaker.py, test_a2_verify.py, test_classifier_accuracy.py, test_agy_tiers.py) successfully with 100% accuracy and zero regressions.
Rebuilt the image, redeployed the Pod, and verified that /visualizer and POST /dashboard/save-annotations endpoints persist data correctly directly to the host filesystem.

Summary by Sourcery

Add a dataset visualizer UI and persistence pipeline, improve async-safe stats/timeline persistence, and extend the circuit breaker and dataset tooling for more robust classifier evaluation.

New Features:

  • Expose a static dataset visualizer HTML UI under /visualizer with annotation storage backed by host-mounted /app/data.
  • Add FastAPI endpoints and static mounts for serving classifier data files and saving human annotations to disk.
  • Introduce multiple scripts for extracting prompts from Langfuse, classifying them with gemma4-26a4b, benchmarking the classifier, and re-running classifications across the dataset.

Bug Fixes:

  • Prevent blocking synchronous disk I/O in async request paths by throttling and offloading stats and timeline JSON writes to background executors.
  • Fix per-tier breaker selection in agy_proxy by routing based on model family name instead of hardcoded index positions.
  • Improve prompt extraction and trivial-test filtering to avoid truncation and misclassification of user prompts.
  • Ensure dual circuit breaker behavior remains backward compatible while exposing independent google and vendor breakers.

Enhancements:

  • Refine the FastAPI lifespan lifecycle to always start/stop the background score push task safely and flush stats/timeline on shutdown.
  • Add comprehensive docstrings to core router, circuit breaker, and streaming helpers to satisfy docstring coverage checks and clarify behavior.
  • Expand circuit breaker tests and verification scripts to cover dual breaker semantics and master compatibility.
  • Tighten dataset visualizer UX with clearer agreement/conflict indicators, reviewed state, and safer label handling.
  • Improve dataset prompt stats and logging output for extraction and benchmarking scripts to aid manual review.

Build:

  • Update the router Containerfile to drop the gitignored best_free_model.json from the image build and include the static visualizer assets.
  • Adjust Dependabot config to monitor the router Containerfile path and simplify commit message settings.

CI:

  • Add a GitHub Actions workflow to run circuit breaker and integration verification tests on pushes and pull requests.

Deployment:

  • Mount a persistent /app/data volume into the router Pod for dataset and annotation storage.

Summary by CodeRabbit

Release Notes

  • New Features

    • Added interactive dataset visualizer UI with filtering and human annotation capabilities
    • Introduced new /visualizer endpoint and annotation storage via /dashboard/save-annotations
    • Made stats persistence asynchronous and non-blocking for improved responsiveness
  • Tests

    • Added automated testing workflow via GitHub Actions
  • Chores

    • Updated infrastructure and test configurations
    • Added utility scripts for dataset processing and model classification

google-labs-jules Bot and others added 19 commits June 18, 2026 10:27
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
- Move _last_stats_save inside try block (only advance on successful write)
- Add save_persisted_stats(force=True) + timeline flush in lifespan shutdown
- Remove redundant import time inside save_persisted_stats
- Fix misleading comment about independent throttle timers
- Use time.monotonic() instead of time.time() for throttle timestamps
- Add logger.warning to bare except in timeline write
- Add logging to bare except Exception: pass in shutdown timeline flush
- Move asyncio.create_task(push_aggregate_scores()) before yield so it runs during app lifetime
- Use atomic file writes (temp file + os.replace) in save_persisted_stats
- Use atomic file writes in shutdown timeline flush and record_tool_usage timeline write
- Restructure lifespan to single yield point (remove else: yield; return pattern)
- Capture time.monotonic() once in record_tool_usage and reuse for guard + assignment
- Extract shared _atomic_write_json_sync / _atomic_write_json_async helpers
- Make save_persisted_stats async, using _atomic_write_json_async with
  run_in_executor + deep-copy to prevent concurrent modification
- Update all 4 async call sites (lifespan, classify_request x2, chat_completions)
  to await save_persisted_stats()
- In record_tool_usage (sync), fire-and-forget both stats and timeline writes
  via loop.run_in_executor with deep-copied data; fall back to sync write
  when no event loop is running (early startup)
- Replace duplicated atomic write logic in lifespan shutdown with shared helper
@sourcery-ai

sourcery-ai Bot commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

Reviewer's Guide

This PR finalizes the classifier dataset visualizer workflow, adds persistent dataset annotation storage, hardens circuit breaker behavior and tests, and optimizes router stats persistence to be async, throttled, and atomic, while also refining data extraction/benchmarking scripts and CI/dependency configs.

Sequence diagram for dataset visualizer and annotation persistence

sequenceDiagram
    actor Reviewer
    participant FastAPI as FastAPI_app
    participant Static as StaticFiles
    participant Browser as Visualizer_JS
    participant SaveAnn as save_annotations
    participant Writer as _atomic_write_json_async

    Reviewer->>FastAPI: GET /visualizer (get_visualizer)
    FastAPI-->>Reviewer: HTMLResponse(visualizer.html)

    Reviewer->>Browser: Open visualizer.html
    Browser->>Static: GET /data/classified_dataset.json
    Static-->>Browser: classified_dataset.json
    Browser->>Static: GET /data/annotations.json
    Static-->>Browser: annotations.json
    Browser->>Browser: init() & render()

    Reviewer->>Browser: Click "Export Annotations"
    Browser->>SaveAnn: POST /dashboard/save-annotations
    SaveAnn->>Writer: _atomic_write_json_async(annotations.json, body)
    Writer-->>SaveAnn: write complete
    SaveAnn-->>Browser: JSONResponse(status="ok")
Loading

Sequence diagram for async, throttled router stats persistence

sequenceDiagram
    participant Chat as chat_completions
    participant Classify as classify_request
    participant Tool as record_tool_usage
    participant SaveStats as save_persisted_stats
    participant AsyncWrite as _atomic_write_json_async
    participant SyncWrite as _atomic_write_json_sync

    Chat->>Classify: classify_request(...)
    Classify->>Classify: update stats["cache_hits"]
    Classify->>SaveStats: await save_persisted_stats()
    SaveStats->>SaveStats: check throttle (_last_stats_save)
    SaveStats->>AsyncWrite: _atomic_write_json_async(STATS_JSON_PATH, stats)
    AsyncWrite->>SyncWrite: _atomic_write_json_sync(path, data_copy)
    SyncWrite-->>AsyncWrite: atomic os.replace
    AsyncWrite-->>SaveStats: done

    Tool->>Tool: update stats and stats["timeline"]
    Tool->>SaveStats: loop.create_task(save_persisted_stats())
    Tool->>SyncWrite: _atomic_write_json_sync(router_timeline.json, timeline) [throttled]
    SyncWrite-->>Tool: done
Loading

File-Level Changes

Change Details Files
Make router stats/timeline persistence non-blocking, atomic, and throttled, and ensure clean shutdown flushing.
  • Introduce synchronous and async atomic JSON writers using temp files and os.replace, with deep-copying before background serialization.
  • Replace save_persisted_stats with an async, throttled implementation that offloads disk writes to a thread pool and is used from classify_request, chat_completions, and record_tool_usage.
  • Start and track the push_aggregate_scores task in lifespan, cancel it on shutdown, and flush both stats and timeline to disk on clean shutdown.
router/main.py
Expose dataset visualizer and static/data directories via FastAPI, with a POST endpoint to persist human annotations to disk.
  • Mount /static and /data directories in the FastAPI app, ensuring a DATA_DIR exists for persisted files.
  • Serve the new visualizer.html via the /visualizer route and link it from the dashboard header.
  • Add a Pydantic AnnotationItem model and a /dashboard/save-annotations POST handler that atomically writes annotations.json to DATA_DIR using async atomic writes.
router/main.py
router/static/visualizer.html
pod.yaml
router/Containerfile
Refine dual circuit breaker behavior and update tests and verification scripts to exercise google/vendor sub-breakers and backward compatibility.
  • Clarify PerModelBreaker and DualCircuitBreaker with docstrings, adjust record_failure logic, and make DualCircuitBreaker.tier return the max sub-tier.
  • Update tests to import router modules via relative paths, add helper reset_breakers, and extend coverage for independent sub-breakers and master compatibility.
  • Adjust verify_breaker and test_a2_verify to validate both google and vendor breakers and ensure master tier reflects sub-breakers.
router/circuit_breaker.py
test_circuit_breaker.py
verify_breaker.py
test_a2_verify.py
Improve prompt extraction and classification tooling for building and benchmarking the routing dataset.
  • Rewrite extract_prompts.py to use the first real user message, apply stricter trivial prompt filtering with regex word boundaries, increase pagination, and log richer stats.
  • Add scripts for gap-fill extraction, complex prompt extraction, direct classification, reclassification with grammar via router, retrying errors with truncation, and benchmarking classifier accuracy and confusion matrix.
  • Ensure scripts read/write from the shared data directory and update classified_dataset.json/benchmark_results.json formats to include llm_tier and clf_tier with agreement stats.
scripts/extract_prompts.py
scripts/extract_gapfill.py
scripts/extract_complex.py
scripts/classify_direct.py
scripts/reclassify_all.py
scripts/retry_errors.py
scripts/benchmark_classifier.py
Configure Dependabot and CI workflow for the router project and add lightweight engineering notes.
  • Scope Docker Dependabot updates to /router, simplify the commit-message configuration, and remove unused ignore rules.
  • Add a GitHub Actions workflow to run the circuit breaker and integration tests on pushes/PRs to master.
  • Add a bolt.md learning note documenting prior synchronous I/O issues in async handlers.
.github/dependabot.yml
.github/workflows/test.yml
.jules/bolt.md

Tips and commands

Interacting with Sourcery

  • Trigger a new review: Comment @sourcery-ai review on the pull request.
  • Continue discussions: Reply directly to Sourcery's review comments.
  • Generate a GitHub issue from a review comment: Ask Sourcery to create an
    issue from a review comment by replying to it. You can also reply to a
    review comment with @sourcery-ai issue to create an issue from it.
  • Generate a pull request title: Write @sourcery-ai anywhere in the pull
    request title to generate a title at any time. You can also comment
    @sourcery-ai title on the pull request to (re-)generate the title at any time.
  • Generate a pull request summary: Write @sourcery-ai summary anywhere in
    the pull request body to generate a PR summary at any time exactly where you
    want it. You can also comment @sourcery-ai summary on the pull request to
    (re-)generate the summary at any time.
  • Generate reviewer's guide: Comment @sourcery-ai guide on the pull
    request to (re-)generate the reviewer's guide at any time.
  • Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
    pull request to resolve all Sourcery comments. Useful if you've already
    addressed all the comments and don't want to see them anymore.
  • Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
    request to dismiss all existing Sourcery reviews. Especially useful if you
    want to start fresh with a new review - don't forget to comment
    @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

  • Enable or disable review features such as the Sourcery-generated pull request
    summary, the reviewer's guide, and others.
  • Change the review language.
  • Add, remove or edit custom review instructions.
  • Adjust other review settings.

Getting Help

@coderabbitai

coderabbitai Bot commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Warning

Review limit reached

@sheepdestroyer, we couldn't start this review because you've reached your PR review rate limit.

More reviews will be available in 49 minutes and 19 seconds. Learn how PR review limits work.

Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file).

⌛ How to resolve this issue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits.

🚦 How do rate limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan refill rate.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, the refill rate gradually slows as usage increases. The highest same-day bursts are limited more strictly.

Please see our Fair Usage Limits Policy for further information.

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 29a29d58-d219-42d8-83d7-b601b330cb75

📥 Commits

Reviewing files that changed from the base of the PR and between e37b24b and 219dd8d.

📒 Files selected for processing (2)
  • router/main.py
  • router/static/visualizer.html
📝 Walkthrough

Walkthrough

This PR refactors the circuit breaker to a DualCircuitBreaker model with separate google and vendor sub-breakers, selected by model name. Stats persistence is converted to an async, throttled, atomic writer. A dataset annotation visualizer HTML page is added with two new backend routes. Seven standalone dataset pipeline scripts are added or updated for prompt extraction, classification, and benchmarking. CI and infra configs are updated.

Changes

Router: DualCircuitBreaker, async persistence, and dataset visualizer

Layer / File(s) Summary
DualCircuitBreaker logic and agy_proxy breaker selection
router/circuit_breaker.py, router/agy_proxy.py
PerModelBreaker.record_failure drops the probe_granted branch; DualCircuitBreaker gains docstrings; agy_proxy selects the google vs vendor sub-breaker based on whether model_name contains "gemini" rather than using tier index.
Dual circuit breaker tests and CI workflow
test_circuit_breaker.py, verify_breaker.py, test_a2_verify.py, .github/workflows/test.yml
test_circuit_breaker.py is rewritten for the dual-breaker model with a reset_breakers() helper, sub-breaker-scoped assertions, and a new test_backward_compatibility(); verify_breaker.py loops over both sub-breakers; test_a2_verify.py applies blocked state to both; a new GitHub Actions workflow runs all three on push/PR.
Atomic write helpers and async save_persisted_stats
router/main.py, .jules/bolt.md
Adds _atomic_write_json_sync and _atomic_write_json_async; converts save_persisted_stats to async def save_persisted_stats(force=False) with throttle tracking; updates lifespan to create/cancel the score-push task and flush stats/timeline atomically on shutdown.
Non-blocking persistence call sites
router/main.py
classify_request, chat_completions await the new async save_persisted_stats(); record_tool_usage uses loop.create_task / run_in_executor with independent timeline throttle and a logged completion callback. Minor docstrings added throughout.
Visualizer backend routes and AnnotationItem model
router/main.py
Adds AnnotationItem Pydantic model; mounts /static and /data as StaticFiles; adds GET /visualizer (serves visualizer.html) and POST /dashboard/save-annotations (persists payload to data/annotations.json atomically); adds a Dataset Visualizer link in the dashboard header.
Visualizer HTML/JS front-end
router/static/visualizer.html
Full dark-themed annotation UI that fetches classified_dataset.json, benchmark_results.json, and annotations.json on load, renders a filterable prompt list with conflict/agree/unreviewed tags, and allows tier selection and note saving with localStorage persistence and POST export.

Dataset pipeline scripts

Layer / File(s) Summary
Prompt extraction scripts
scripts/extract_prompts.py, scripts/extract_complex.py, scripts/extract_gapfill.py
extract_prompts.py switches to extract_first_user_prompt with regex-based trivial filtering and writes raw_prompts_v2.json; extract_complex.py collects complex prompts from Langfuse using keyword heuristics; extract_gapfill.py fetches DEFAULT-level observations with length-biased deduplication.
Classification, reclassification, retry, and benchmark
scripts/classify_direct.py, scripts/reclassify_all.py, scripts/retry_errors.py, scripts/benchmark_classifier.py
classify_direct.py classifies raw Hermes prompts and writes classified_dataset.json; reclassify_all.py re-runs the classifier for both llm_tier and clf_tier with atomic save; retry_errors.py discovers model port and retries ERROR entries in batches; benchmark_classifier.py evaluates tier accuracy with a confusion matrix and writes benchmark_results.json.

Infrastructure: container, deployment, and Dependabot

Layer / File(s) Summary
Container image, pod volume, and Dependabot config
router/Containerfile, pod.yaml, .github/dependabot.yml
Containerfile stages free_models_roster.json instead of best_free_model.json; pod.yaml mounts a dataset-data hostPath volume at /app/data; Dependabot Docker directory updated to /router with prefix-major override removed.

Sequence Diagram(s)

sequenceDiagram
    participant Handler as API Handler (classify_request / chat_completions)
    participant SPS as save_persisted_stats (async)
    participant AWA as _atomic_write_json_async
    participant Executor as ThreadPoolExecutor
    participant Disk as stats JSON (disk)

    Handler->>SPS: await save_persisted_stats()
    SPS->>SPS: check throttle (_last_stats_save)
    SPS->>AWA: deepcopy stats, await _atomic_write_json_async(path, data)
    AWA->>Executor: run_in_executor(_atomic_write_json_sync)
    Executor->>Disk: write to tempfile, os.replace
    Disk-->>Executor: done
    Executor-->>AWA: future result
    AWA-->>SPS: complete (exception logged on failure)
    SPS-->>Handler: return
Loading
sequenceDiagram
    participant Browser as Browser
    participant Visualizer as GET /visualizer
    participant SaveAnn as POST /dashboard/save-annotations
    participant AWA as _atomic_write_json_async
    participant DataDir as data/annotations.json

    Browser->>Visualizer: GET /visualizer
    Visualizer-->>Browser: visualizer.html (or 404)
    Browser->>Browser: fetch /data/classified_dataset.json
    Browser->>Browser: fetch /data/annotations.json (prefill)
    Browser->>SaveAnn: POST /dashboard/save-annotations (AnnotationItem dict)
    SaveAnn->>AWA: await _atomic_write_json_async(annotations.json, payload)
    AWA->>DataDir: atomic write
    DataDir-->>SaveAnn: done
    SaveAnn-->>Browser: {"status": "ok"}
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Poem

🐇 Hoppity-hop through the async lane,
No more blocking on the I/O train!
Dual breakers guard both google and friend,
Annotations saved before the end.
The visualizer blooms on the screen,
The tidiest dataset there has been! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name Status Explanation Resolution
Title check ❓ Inconclusive The title 'Finalize pr2 fixes review' is vague and generic, using non-descriptive terms that don't convey meaningful information about the changeset. Consider a more specific title that highlights the main changes, such as 'Fix UI state sync, improve breaker mapping, and add docstrings' or reference the primary focus from the objectives.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage ✅ Passed Docstring coverage is 98.11% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch finalize-pr2-fixes-review

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey - I've found 1 issue

Prompt for AI Agents
Please address the comments from this code review:

## Individual Comments

### Comment 1
<location path="router/main.py" line_range="2565-2574" />
<code_context>
+    ts: Optional[str] = None
+
+@app.post("/dashboard/save-annotations")
+async def save_annotations(payload: Dict[str, AnnotationItem]):
+    """Save human review annotations to disk."""
+    try:
+        body = {k: (v.model_dump() if hasattr(v, "model_dump") else v.dict()) for k, v in payload.items()}
+        ann_path = DATA_DIR / "annotations.json"
+        await _atomic_write_json_async(str(ann_path), body)
+        return JSONResponse({"status": "ok", "saved": len(body)})
+    except Exception as e:
+        logger.error(f"Failed to save annotations: {e}")
+        raise HTTPException(status_code=500, detail="Failed to save annotations")
+
 if __name__ == "__main__":
</code_context>
<issue_to_address>
**🚨 issue (security):** Annotations save endpoint is unauthenticated and could be abused or corrupted by arbitrary clients

This handler accepts arbitrary JSON from any caller and writes it straight to disk. If this endpoint is exposed beyond a trusted admin network, unauthenticated users can overwrite or spam your annotation data. Please protect it with the same auth/secret mechanism as the dashboard (or a shared header/token) before allowing writes. Also consider validating that keys are numeric indices and enforcing a max payload size to avoid abuse with large or malformed inputs.
</issue_to_address>

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

Comment thread router/main.py
Comment on lines +2565 to +2574
async def save_annotations(payload: Dict[str, AnnotationItem]):
"""Save human review annotations to disk."""
try:
body = {k: (v.model_dump() if hasattr(v, "model_dump") else v.dict()) for k, v in payload.items()}
ann_path = DATA_DIR / "annotations.json"
await _atomic_write_json_async(str(ann_path), body)
return JSONResponse({"status": "ok", "saved": len(body)})
except Exception as e:
logger.error(f"Failed to save annotations: {e}")
raise HTTPException(status_code=500, detail="Failed to save annotations")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚨 issue (security): Annotations save endpoint is unauthenticated and could be abused or corrupted by arbitrary clients

This handler accepts arbitrary JSON from any caller and writes it straight to disk. If this endpoint is exposed beyond a trusted admin network, unauthenticated users can overwrite or spam your annotation data. Please protect it with the same auth/secret mechanism as the dashboard (or a shared header/token) before allowing writes. Also consider validating that keys are numeric indices and enforcing a max payload size to avoid abuse with large or malformed inputs.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a dual circuit breaker system separating Google and vendor models, refactors telemetry persistence to use non-blocking atomic writes, and adds a dataset visualizer tool along with several classifier benchmarking and prompt extraction scripts. The review feedback highlights two key areas for improvement: first, automatically persisting human annotations to the server in the visualizer to prevent data loss, and second, addressing dead or unreachable streaming logic within the proxy implementation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment on lines +328 to +331
function saveAnnotations() {
// Save to localStorage and trigger download
localStorage.setItem('classifier-annotations', JSON.stringify(annotations));
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The saveAnnotations function currently only persists annotations to the browser's localStorage. This means any human review progress is not automatically saved to the server's disk (/app/data/annotations.json) until the user manually clicks the Export Annotations button.

If the user closes the tab, reloads, or if another reviewer accesses the visualizer, the state will be out of sync.

To ensure seamless persistence and prevent data loss, consider automatically posting the updated annotations to the /dashboard/save-annotations endpoint whenever saveAnnotations is called.

Suggested change
function saveAnnotations() {
// Save to localStorage and trigger download
localStorage.setItem('classifier-annotations', JSON.stringify(annotations));
}
function saveAnnotations() {
localStorage.setItem('classifier-annotations', JSON.stringify(annotations));
fetch('/dashboard/save-annotations', {
method: 'POST',
headers: {'Content-Type': 'application/json'},
body: JSON.stringify(annotations),
}).catch(err => console.warn('Failed to auto-save annotations to server:', err));
}

Comment thread router/agy_proxy.py
Comment on lines 326 to +327
async def token_generator(stream_resp, httpx_client, initial_line, current_conv_id):
"""Asynchronously yields tokens from the agy daemon stream and manages session state updates."""

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The stream=True parameter and the token_generator async generator in try_agy_proxy appear to be dead/unreachable code.

In router/main.py, try_agy_proxy is always called with the default stream=False (non-streaming). If a client requests a stream, main.py still awaits the full synchronous response from try_agy_proxy and then simulates a stream post-hoc using agy_stream_generator.

Furthermore, if stream=True were actually passed to try_agy_proxy, main.py would crash on lines like agy_response.get("choices") because try_agy_proxy returns a dictionary containing {"stream": ..., "model": ...} instead of the expected OpenAI-compatible structure.

Consider either fully integrating the real stream generator in main.py or removing this unused streaming logic from agy_proxy.py to simplify maintenance.

Comment thread router/main.py Dismissed
Comment thread router/main.py Dismissed
Comment thread router/main.py Dismissed

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 10

🧹 Nitpick comments (4)
pod.yaml (2)

471-472: 💤 Low value

Awareness: Hardcoded user-specific path limits portability.

The hostPath uses /home/gpav/Vrac/LAB/AI/LLM-Routing/data, which is specific to the development environment. This is acceptable for local development but would need to be parameterized or templated for broader deployment (e.g., via Kustomize overlays or Helm values).

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pod.yaml` around lines 471 - 472, The hostPath volume definition uses a
hardcoded user-specific path that reduces portability across different
environments and machines. Replace the hardcoded path
`/home/gpav/Vrac/LAB/AI/LLM-Routing/data` with a parameterized value by either
creating a Kustomize overlay that patches the path for different environments,
using Helm chart values with a configurable data path parameter, or using
environment variable substitution in your deployment pipeline. This allows the
pod.yaml to be reused across local development, staging, and production
environments without modification.

1-474: Operational advice: Consider addressing container security posture.

Static analysis identified three security concerns that apply to the entire manifest (pre-existing, not introduced by this PR):

  1. allowPrivilegeEscalation: Containers should explicitly set securityContext.allowPrivilegeEscalation: false to prevent privilege escalation.
  2. Root containers: Consider running containers with non-root users via securityContext.runAsNonRoot: true and explicit runAsUser values.
  3. Host network: The pod uses hostNetwork: true, which shares the host's network namespace and bypasses network policies.

These are operational considerations for a production hardening review rather than blockers for this PR's dataset volume mount.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pod.yaml` around lines 1 - 474, Add security context configurations to the
pod specification to address three security concerns. First, add a
`securityContext` section at the pod spec level (under spec) to define default
security settings for all containers, or add individual `securityContext`
sections to each container definition. Set `allowPrivilegeEscalation: false` to
prevent privilege escalation, add `runAsNonRoot: true` and specify an explicit
`runAsUser` value (e.g., 1000) to run containers with non-root users instead of
root. Additionally, review the necessity of `hostNetwork: true` at the pod spec
level, as it bypasses network policies by sharing the host's network namespace,
and consider removing or restricting it if not absolutely required for the
application functionality.
verify_breaker.py (1)

11-11: Replace == True/False assertions with direct boolean assertions to align with PEP 8 and resolve Ruff E712.

Normalize assertions at lines 11, 14, 17, 20, and 21 by using assert b.is_allowed() instead of assert b.is_allowed() == True and assert not sub.is_allowed() instead of assert sub.is_allowed() == False.

♻️ Suggested cleanup
-assert b.is_allowed() == True, 'Tier 0 should be open'
+assert b.is_allowed(), 'Tier 0 should be open'
...
-    assert sub.is_allowed() == False, 'Tier 1 should block (cooldown active)'
+    assert not sub.is_allowed(), 'Tier 1 should block (cooldown active)'
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@verify_breaker.py` at line 11, The assertions in verify_breaker.py are using
explicit `== True` and `== False` comparisons which violates PEP 8 and Ruff E712
linting rules. Replace the assertion at line 11 that compares `b.is_allowed() ==
True` with a direct boolean assertion `assert b.is_allowed()`, and similarly
replace all assertions comparing `sub.is_allowed() == False` at lines 14, 17,
20, and 21 with `assert not sub.is_allowed()` to use direct truthiness checks
instead of explicit boolean comparisons.

Source: Linters/SAST tools

test_circuit_breaker.py (1)

38-41: Apply idiomatic truthiness assertions to remove Ruff E712 violations.

Replace == True/False comparisons with direct assertions. Use assert expr for True checks and assert not expr for False checks across all 14 instances.

♻️ Refactoring pattern
-    assert b.is_allowed() == True
+    assert b.is_allowed()
...
-    assert b.google.is_allowed() == False
+    assert not b.google.is_allowed()

Lines: 38, 40, 41, 53, 56, 69, 70, 71, 86, 103, 118, 130, 135

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@test_circuit_breaker.py` around lines 38 - 41, Replace all explicit
True/False comparisons in assertions with idiomatic truthiness checks to resolve
Ruff E712 violations. For assertions checking True values like is_allowed(),
google.is_allowed(), and vendor.is_allowed(), remove the == True comparison and
use a direct assert statement. For any assertions checking False values (if
present), use assert not instead. Apply this pattern consistently across all 14
instances mentioned in the comment (lines 38, 40, 41, 53, 56, 69, 70, 71, 86,
103, 118, 130, 135).

Source: Linters/SAST tools

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In @.github/workflows/test.yml:
- Around line 13-19: Replace the floating version tags (`@v4` and `@v5`) in the
GitHub Actions workflow with specific commit SHAs for both the Checkout
Repository step (actions/checkout) and the Set up Python step
(actions/setup-python). Additionally, add persist-credentials: false as a
configuration option under the Checkout Repository step to prevent credential
exposure in PR workflows. These changes pin the exact action versions to prevent
supply-chain vulnerabilities from action re-releases.

In `@router/main.py`:
- Around line 2547-2570: The static file mount for "/data" and the
save_annotations endpoint lack authentication, exposing them to unauthorized
access. Add route-level authentication to both the app.mount("/data", ...) call
and the save_annotations function by implementing a Depends() dependency that
validates API keys or authorization tokens. Consider creating a reusable
authentication dependency function that validates incoming requests and returns
an authenticated user or raises an HTTPException, then apply this dependency to
both the data mount configuration and as a parameter to the save_annotations
endpoint to ensure only authenticated requests can read or write data.
- Around line 706-711: The create_task call on line 710 does not retain a
reference to the returned task, which can cause the background stats-save
operation to be garbage collected before completion, violating Ruff RUF006.
Create a module-level set to store pending background tasks, then after calling
loop.create_task(save_persisted_stats()) assign the returned task to this set
and attach a done callback that removes the task from the set when it completes,
ensuring the fire-and-forget telemetry write is not lost.
- Around line 393-408: The cleanup code for persisting stats and timeline (the
calls to save_persisted_stats and _atomic_write_json_async) may not execute if
an exception occurs during the yield statement. Wrap the yield and the task
cancellation logic in a try/finally block, with the stats/timeline persistence
code in the finally block to guarantee cleanup always runs regardless of
exceptions during the main execution path.

In `@router/static/visualizer.html`:
- Around line 157-158: Replace the array index-based keying mechanism for
annotations with stable identifiers such as a prompt hash, session_id, or
explicit id throughout the annotation system. In the getAnnotation function and
related annotation storage/retrieval code (lines 307-324, 330, 346-347), update
the key structure to use these stable identifiers instead of array indices.
Maintain backward compatibility by adding a fallback mechanism that attempts to
resolve annotations by the legacy index-based key if the stable identifier
lookup fails, allowing existing saved annotations to continue working during
migration.
- Around line 228-232: The div element with class "list-item" that renders
selectable prompt rows is currently only mouse-operable via the onclick handler
for selectPrompt. To make this keyboard-operable for accessibility, convert the
div element to a button element or add keyboard event handling and accessibility
attributes (tabindex, role="button", and onkeydown handler) to the div. This
will allow keyboard-only users to focus on and select prompts using Tab and
Enter keys, enabling them to use the core review workflow.

In `@scripts/benchmark_classifier.py`:
- Around line 53-56: The benchmark script assumes every item in the prompts
dataset has a "tier" key and that the value is always one of the valid TIERS,
but datasets from reclassify_all.py use "llm_tier" or "clf_tier" instead, and
may contain unresolved labels like "ERROR". Replace the direct access to
item["tier"] with a safe lookup using .get() that checks for multiple possible
key names (tier, llm_tier, clf_tier), and add defensive validation around line
68 where the tier value is checked against TIERS to handle unexpected values
gracefully rather than raising a KeyError that crashes the benchmark run.

In `@scripts/extract_complex.py`:
- Around line 40-55: The extract_user_prompt function currently iterates through
messages in reversed order using reversed(messages), which causes it to select
the last user message instead of the first. Change the loop to iterate forward
through messages in their original order by removing the reversed() call and
updating for msg in reversed(messages) to for msg in messages. This will ensure
the function returns the first user message, aligning with the behavior already
implemented in scripts/extract_prompts.py and preventing the prompt-selection
issue mentioned in the review.

In `@scripts/extract_gapfill.py`:
- Around line 38-56: The extract_user_prompt function must use the same prompt
extraction semantics as the primary extractor to avoid reintroducing the bad
prompt capture issue fixed elsewhere. Currently, the function walks through
messages in reversed order looking for user role entries without filtering
pseudo-user system notes. Locate the primary extractor's implementation in the
codebase and apply its same message filtering and validation logic to the
message extraction loop in extract_user_prompt, ensuring both functions handle
user prompts consistently.

In `@scripts/retry_errors.py`:
- Around line 70-71: The retry logic needs to be schema-aware to handle both the
old tier schema (single `tier` field) and the new schema used by
reclassify_all.py (which uses `llm_tier` and `clf_tier`). Update the list
comprehension at line 70 that filters for `tier == 'ERROR'` to check for both
old and new schema fields, then apply the same schema-aware logic consistently
throughout lines 82-98 when reading and counting tier values (currently line 95
references `llm_tier` inconsistently). Additionally, the `fixed` counter should
only be incremented when the classify() function returns a successful result
(not "ERROR"), so add a check before incrementing the counter to ensure the
classification actually fixed the error.

---

Nitpick comments:
In `@pod.yaml`:
- Around line 471-472: The hostPath volume definition uses a hardcoded
user-specific path that reduces portability across different environments and
machines. Replace the hardcoded path `/home/gpav/Vrac/LAB/AI/LLM-Routing/data`
with a parameterized value by either creating a Kustomize overlay that patches
the path for different environments, using Helm chart values with a configurable
data path parameter, or using environment variable substitution in your
deployment pipeline. This allows the pod.yaml to be reused across local
development, staging, and production environments without modification.
- Around line 1-474: Add security context configurations to the pod
specification to address three security concerns. First, add a `securityContext`
section at the pod spec level (under spec) to define default security settings
for all containers, or add individual `securityContext` sections to each
container definition. Set `allowPrivilegeEscalation: false` to prevent privilege
escalation, add `runAsNonRoot: true` and specify an explicit `runAsUser` value
(e.g., 1000) to run containers with non-root users instead of root.
Additionally, review the necessity of `hostNetwork: true` at the pod spec level,
as it bypasses network policies by sharing the host's network namespace, and
consider removing or restricting it if not absolutely required for the
application functionality.

In `@test_circuit_breaker.py`:
- Around line 38-41: Replace all explicit True/False comparisons in assertions
with idiomatic truthiness checks to resolve Ruff E712 violations. For assertions
checking True values like is_allowed(), google.is_allowed(), and
vendor.is_allowed(), remove the == True comparison and use a direct assert
statement. For any assertions checking False values (if present), use assert not
instead. Apply this pattern consistently across all 14 instances mentioned in
the comment (lines 38, 40, 41, 53, 56, 69, 70, 71, 86, 103, 118, 130, 135).

In `@verify_breaker.py`:
- Line 11: The assertions in verify_breaker.py are using explicit `== True` and
`== False` comparisons which violates PEP 8 and Ruff E712 linting rules. Replace
the assertion at line 11 that compares `b.is_allowed() == True` with a direct
boolean assertion `assert b.is_allowed()`, and similarly replace all assertions
comparing `sub.is_allowed() == False` at lines 14, 17, 20, and 21 with `assert
not sub.is_allowed()` to use direct truthiness checks instead of explicit
boolean comparisons.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 7792a97e-60f4-493e-8145-f646509255eb

📥 Commits

Reviewing files that changed from the base of the PR and between 859fb74 and e37b24b.

📒 Files selected for processing (19)
  • .github/dependabot.yml
  • .github/workflows/test.yml
  • .jules/bolt.md
  • pod.yaml
  • router/Containerfile
  • router/agy_proxy.py
  • router/circuit_breaker.py
  • router/main.py
  • router/static/visualizer.html
  • scripts/benchmark_classifier.py
  • scripts/classify_direct.py
  • scripts/extract_complex.py
  • scripts/extract_gapfill.py
  • scripts/extract_prompts.py
  • scripts/reclassify_all.py
  • scripts/retry_errors.py
  • test_a2_verify.py
  • test_circuit_breaker.py
  • verify_breaker.py

Comment on lines +13 to +19
- name: Checkout Repository
uses: actions/checkout@v4

- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.11'

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
rg -n 'uses:\s*actions/(checkout|setup-python)@' .github/workflows/test.yml
rg -n 'persist-credentials' .github/workflows/test.yml || true

Repository: sheepdestroyer/LLM-Routing

Length of output: 148


Pin GitHub Actions by commit SHA and disable checkout credential persistence.

Using floating action tags (@v4, @v5) instead of commit SHAs creates a supply-chain security vulnerability—the underlying action can change when re-released. Additionally, persist-credentials: false must be set on the checkout action to prevent credential exposure in PR workflows.

🔒 Proposed workflow hardening
      - name: Checkout Repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@<full-length-commit-sha>
+        with:
+          persist-credentials: false

      - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@<full-length-commit-sha>
🧰 Tools
🪛 zizmor (1.25.2)

[warning] 13-14: credential persistence through GitHub Actions artifacts (artipacked): does not set persist-credentials: false

(artipacked)


[error] 14-14: unpinned action reference (unpinned-uses): action is not pinned to a hash (required by blanket policy)

(unpinned-uses)


[error] 17-17: unpinned action reference (unpinned-uses): action is not pinned to a hash (required by blanket policy)

(unpinned-uses)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In @.github/workflows/test.yml around lines 13 - 19, Replace the floating
version tags (`@v4` and `@v5`) in the GitHub Actions workflow with specific commit
SHAs for both the Checkout Repository step (actions/checkout) and the Set up
Python step (actions/setup-python). Additionally, add persist-credentials: false
as a configuration option under the Checkout Repository step to prevent
credential exposure in PR workflows. These changes pin the exact action versions
to prevent supply-chain vulnerabilities from action re-releases.

Source: Linters/SAST tools

Comment thread router/main.py Outdated
Comment on lines +393 to +408
yield

# Cancel background score task
task.cancel()
try:
await sync_adaptive_router_roster(litellm_master_key)
await task
except asyncio.CancelledError:
pass

# Flush any buffered stats/timeline on clean shutdown
await save_persisted_stats(force=True)
try:
timeline_path = os.path.join(os.path.dirname(CONFIG_PATH), "router_timeline.json")
await _atomic_write_json_async(timeline_path, stats["timeline"])
except Exception as e:
logger.error(f"Roster sync failed: {e}")
yield
# Start aggregate score-push background task (runs for server lifetime)
asyncio.create_task(push_aggregate_scores())
logger.warning(f"Failed to persist timeline on shutdown: {e}")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

cd router && head -n 420 main.py | tail -n 50

Repository: sheepdestroyer/LLM-Routing

Length of output: 1800


🏁 Script executed:

cd router && head -n 393 main.py | tail -n 50 | head -n 30

Repository: sheepdestroyer/LLM-Routing

Length of output: 1596


🏁 Script executed:

python3 <<'PY'
import asyncio
from contextlib import asynccontextmanager

# Test 1: Does code after yield run if an exception is raised in the context?
print("=== Test 1: Exception in context ===")

`@asynccontextmanager`
async def test_context_1():
    print("Setup")
    yield
    print("Cleanup after yield")

async def test_1():
    try:
        async with test_context_1():
            print("Inside context")
            raise ValueError("Test error")
    except ValueError:
        print("Caught exception")

asyncio.run(test_1())

# Test 2: Does code after yield run if await raises an exception?
print("\n=== Test 2: Exception from await in cleanup ===")

`@asynccontextmanager`
async def test_context_2():
    async def failing_task():
        await asyncio.sleep(0.1)
        raise ValueError("Task failed")
    
    task = asyncio.create_task(failing_task())
    print("Setup")
    yield
    print("Cleanup: about to await task")
    try:
        await task
    except ValueError as e:
        print(f"Caught task exception: {e}")
    print("Cleanup: after await")

async def test_2():
    async with test_context_2():
        print("Inside context")

asyncio.run(test_2())

# Test 3: What if we don't catch the exception from await?
print("\n=== Test 3: Uncaught exception from await ===")

`@asynccontextmanager`
async def test_context_3():
    async def failing_task():
        await asyncio.sleep(0.1)
        raise ValueError("Task failed")
    
    task = asyncio.create_task(failing_task())
    print("Setup")
    yield
    print("Cleanup: about to await task")
    await task  # This will raise!
    print("Cleanup: after await (will not execute)")

async def test_3():
    try:
        async with test_context_3():
            print("Inside context")
    except ValueError as e:
        print(f"Exception propagated: {e}")

asyncio.run(test_3())
PY

Repository: sheepdestroyer/LLM-Routing

Length of output: 431


🏁 Script executed:

python3 <<'PY'
import asyncio

# Test: What happens when we cancel and await a task?
print("=== Test: await on cancelled task ===")

async def test_task():
    try:
        await asyncio.sleep(10)
    except asyncio.CancelledError:
        print("Task received CancelledError")
        raise  # Proper cleanup behavior

async def main():
    task = asyncio.create_task(test_task())
    await asyncio.sleep(0.01)  # Let it start
    task.cancel()
    
    try:
        await task
    except asyncio.CancelledError:
        print("Caught CancelledError from await")
    except Exception as e:
        print(f"Caught other exception: {type(e).__name__}: {e}")

asyncio.run(main())

# Test: What if task raises an exception before we cancel?
print("\n=== Test: task fails before cancel ===")

async def failing_task():
    await asyncio.sleep(0.01)
    raise RuntimeError("Task failed")

async def main2():
    task = asyncio.create_task(failing_task())
    await asyncio.sleep(0.05)  # Let it fail
    
    print(f"Task done: {task.done()}, exception: {task.exception() if task.done() else 'N/A'}")
    
    task.cancel()  # Try to cancel after it's already failed
    
    try:
        await task
    except asyncio.CancelledError:
        print("Caught CancelledError from await")
    except Exception as e:
        print(f"Caught other exception: {type(e).__name__}: {e}")

asyncio.run(main2())
PY

Repository: sheepdestroyer/LLM-Routing

Length of output: 300


🏁 Script executed:

cd router && rg -n "def push_aggregate_scores|async def push_aggregate_scores" -A 15

Repository: sheepdestroyer/LLM-Routing

Length of output: 939


🏁 Script executed:

cd router && sed -n '50,100p' main.py

Repository: sheepdestroyer/LLM-Routing

Length of output: 2691


🏁 Script executed:

python3 <<'PY'
# Verify FastAPI lifespan exception handling
import asyncio
from contextlib import asynccontextmanager
from typing import AsyncGenerator

# Simulate FastAPI's lifespan behavior
`@asynccontextmanager`
async def lifespan_simulation(app_name: str) -> AsyncGenerator:
    """Simulates how FastAPI handles lifespan context managers."""
    print(f"[{app_name}] Startup")
    try:
        yield
    except Exception as e:
        print(f"[{app_name}] Exception during app execution: {e}")
        raise
    finally:
        print(f"[{app_name}] Shutdown/cleanup running")

# Test: What if an exception is raised during app runtime (inside the async with)?
print("=== Test: Exception during app lifetime ===")

async def test_with_exception():
    async with lifespan_simulation("app"):
        print("App running...")
        raise RuntimeError("App crashed!")

try:
    asyncio.run(test_with_exception())
except RuntimeError as e:
    print(f"Exception caught outside: {e}")

# Test 2: Normal shutdown
print("\n=== Test: Normal shutdown ===")

async def test_normal():
    async with lifespan_simulation("app"):
        print("App running...")

asyncio.run(test_normal())
PY

Repository: sheepdestroyer/LLM-Routing

Length of output: 355


Guarantee cleanup with try/finally to ensure stats/timeline persist on all shutdown paths.

If push_aggregate_scores() fails, await task will re-raise the exception and skip the stats/timeline flush. Wrap the yield in try/finally to guarantee cleanup always runs.

Proposed fix
-    yield
-
-    # Cancel background score task
-    task.cancel()
-    try:
-        await task
-    except asyncio.CancelledError:
-        pass
-
-    # Flush any buffered stats/timeline on clean shutdown
-    await save_persisted_stats(force=True)
-    try:
-        timeline_path = os.path.join(os.path.dirname(CONFIG_PATH), "router_timeline.json")
-        await _atomic_write_json_async(timeline_path, stats["timeline"])
-    except Exception as e:
-        logger.warning(f"Failed to persist timeline on shutdown: {e}")
+    try:
+        yield
+    finally:
+        # Cancel background score task
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            logger.warning(f"Background score task failed before shutdown: {e}")
+
+        # Flush any buffered stats/timeline on shutdown
+        await save_persisted_stats(force=True)
+        try:
+            timeline_path = os.path.join(os.path.dirname(CONFIG_PATH), "router_timeline.json")
+            await _atomic_write_json_async(timeline_path, stats["timeline"])
+        except Exception as e:
+            logger.warning(f"Failed to persist timeline on shutdown: {e}")
🧰 Tools
🪛 Ruff (0.15.17)

[warning] 407-407: Do not catch blind exception: Exception

(BLE001)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@router/main.py` around lines 393 - 408, The cleanup code for persisting stats
and timeline (the calls to save_persisted_stats and _atomic_write_json_async)
may not execute if an exception occurs during the yield statement. Wrap the
yield and the task cancellation logic in a try/finally block, with the
stats/timeline persistence code in the finally block to guarantee cleanup always
runs regardless of exceptions during the main execution path.

Comment thread router/main.py Outdated
Comment on lines +706 to +711
# Fire-and-forget stats write via save_persisted_stats (non-blocking)
now = time.monotonic()
try:
timeline_path = os.path.join(os.path.dirname(CONFIG_PATH), "router_timeline.json")
with open(timeline_path, "w") as f:
json.dump(stats["timeline"], f)
except Exception:
pass # disk persistence failure is non-fatal
loop = asyncio.get_running_loop()
loop.create_task(save_persisted_stats())
except RuntimeError:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

# First, check the file exists and read the relevant section
if [ -f "router/main.py" ]; then
  echo "=== Lines 700-720 of router/main.py ==="
  sed -n '700,720p' router/main.py | cat -n
else
  echo "File not found. Let me search for it."
  find . -name "main.py" -path "*/router/*" 2>/dev/null | head -5
fi

Repository: sheepdestroyer/LLM-Routing

Length of output: 1015


🏁 Script executed:

# Also search for create_task patterns and the function name
rg "create_task|save_persisted_stats" --context=3 -n

Repository: sheepdestroyer/LLM-Routing

Length of output: 3362


🏁 Script executed:

# Check if _background_stats_save_tasks already exists in the file
rg "_background_stats_save_tasks" router/main.py

# Also check the exact context around line 193 to see what module-level vars exist
sed -n '190,200p' router/main.py | cat -n

Repository: sheepdestroyer/LLM-Routing

Length of output: 547


Keep a strong reference to the background stats-save task.

Line 710 drops the task returned by create_task(), matching Ruff RUF006; keep it in a module-level set until completion so fire-and-forget telemetry writes are not lost.

Proposed fix
 _last_stats_save = 0.0
+_background_stats_save_tasks = set()
     try:
         loop = asyncio.get_running_loop()
-        loop.create_task(save_persisted_stats())
+        task = loop.create_task(save_persisted_stats())
+        _background_stats_save_tasks.add(task)
+        task.add_done_callback(_background_stats_save_tasks.discard)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# Fire-and-forget stats write via save_persisted_stats (non-blocking)
now = time.monotonic()
try:
timeline_path = os.path.join(os.path.dirname(CONFIG_PATH), "router_timeline.json")
with open(timeline_path, "w") as f:
json.dump(stats["timeline"], f)
except Exception:
pass # disk persistence failure is non-fatal
loop = asyncio.get_running_loop()
loop.create_task(save_persisted_stats())
except RuntimeError:
# Fire-and-forget stats write via save_persisted_stats (non-blocking)
now = time.monotonic()
try:
loop = asyncio.get_running_loop()
task = loop.create_task(save_persisted_stats())
_background_stats_save_tasks.add(task)
task.add_done_callback(_background_stats_save_tasks.discard)
except RuntimeError:
🧰 Tools
🪛 Ruff (0.15.17)

[warning] 710-710: Store a reference to the return value of loop.create_task

(RUF006)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@router/main.py` around lines 706 - 711, The create_task call on line 710 does
not retain a reference to the returned task, which can cause the background
stats-save operation to be garbage collected before completion, violating Ruff
RUF006. Create a module-level set to store pending background tasks, then after
calling loop.create_task(save_persisted_stats()) assign the returned task to
this set and attach a done callback that removes the task from the set when it
completes, ensuring the fire-and-forget telemetry write is not lost.

Source: Linters/SAST tools

Comment thread router/main.py
Comment on lines +2547 to +2570
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
app.mount("/data", StaticFiles(directory=str(DATA_DIR)), name="data")

@app.get("/visualizer", response_class=HTMLResponse)
async def get_visualizer():
"""Serve the dataset visualizer for human review."""
vis_path = STATIC_DIR / "visualizer.html"
if vis_path.exists():
return HTMLResponse(vis_path.read_text())
return HTMLResponse("<h2>Visualizer not found</h2>", status_code=404)

class AnnotationItem(BaseModel):
"""Pydantic model representing a single human dataset review annotation."""
tier: Union[int, str, None] = None
note: str = ""
ts: Optional[str] = None

@app.post("/dashboard/save-annotations")
async def save_annotations(payload: Dict[str, AnnotationItem]):
"""Save human review annotations to disk."""
try:
body = {k: (v.model_dump() if hasattr(v, "model_dump") else v.dict()) for k, v in payload.items()}
ann_path = DATA_DIR / "annotations.json"
await _atomic_write_json_async(str(ann_path), body)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Description: Inspect whether the visualizer/data routes are protected and how the router is exposed.
# Expected: Either auth/dependency/middleware protects these routes, or deployment binds them only to a trusted boundary.

rg -n -C3 'app\.mount\("/data"|`@app`\.(get|post)\("/(visualizer|dashboard/save-annotations)"|Depends\(|APIKey|Authorization|middleware' router/main.py

rg -n -C3 'hostPort|containerPort|0\.0\.0\.0|5000|Ingress|Route|Service|NodePort|LoadBalancer' --iglob '*.yml' --iglob '*.yaml' --iglob 'Containerfile'

Repository: sheepdestroyer/LLM-Routing

Length of output: 4510


🏁 Script executed:

# Check for middleware registration, app-level depends, and other auth mechanisms
rg -n 'app\.middleware|app\.add_middleware|app\.dependency_overrides|APIRouter|Depends\(' router/main.py | head -30

# Check if there are any decorators or auth on the routes specifically
rg -B5 '`@app`\.(get|post|mount)\("/(visualizer|dashboard|data)' router/main.py

# Check for FastAPI app initialization and configuration
rg -n 'FastAPI\(|app\s*=' router/main.py | head -10

Repository: sheepdestroyer/LLM-Routing

Length of output: 935


🏁 Script executed:

# Look at app initialization context and any middleware setup around line 410
sed -n '405,450p' router/main.py

# Search for middleware additions comprehensively
rg -n 'add_middleware|`@app`\.middleware|middleware=' router/main.py

# Check if there's any global dependency or exception handler setup
rg -n '`@app`\.exception_handler|app\.dependency_overrides|__init__' router/main.py

Repository: sheepdestroyer/LLM-Routing

Length of output: 2146


🏁 Script executed:

# Check around the problematic routes to see if there's any Depends() on them
sed -n '2540,2580p' router/main.py

# Search more broadly for any auth/security setup in the entire file
rg -n 'HTTPBearer|HTTPBasic|api_key|check_auth|get_current|security' router/main.py

# Check the top of the file for imports related to security/auth
head -50 router/main.py | grep -E 'import|from'

Repository: sheepdestroyer/LLM-Routing

Length of output: 2761


Add authentication to data routes before production exposure.

The /data mount and /dashboard/save-annotations endpoint have no authentication or authorization. With the app binding to 0.0.0.0:5000 (confirmed in pod.yaml, config.yaml, Containerfile), these routes are externally accessible:

  • /data allows unauthenticated read access to the entire data directory
  • /dashboard/save-annotations allows unauthenticated write access to disk

Either add route-level authentication (e.g., Depends() with API key validation), apply middleware protecting these endpoints, or restrict network exposure to trusted clients only.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@router/main.py` around lines 2547 - 2570, The static file mount for "/data"
and the save_annotations endpoint lack authentication, exposing them to
unauthorized access. Add route-level authentication to both the
app.mount("/data", ...) call and the save_annotations function by implementing a
Depends() dependency that validates API keys or authorization tokens. Consider
creating a reusable authentication dependency function that validates incoming
requests and returns an authenticated user or raises an HTTPException, then
apply this dependency to both the data mount configuration and as a parameter to
the save_annotations endpoint to ensure only authenticated requests can read or
write data.

Comment thread router/static/visualizer.html Outdated
Comment on lines +157 to +158
function getAnnotation(idx) {
return annotations[idx] || null;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Use stable prompt identifiers for annotation keys to avoid silent mislabeling.

Line 157 and Lines 307-324 key annotations by array index. If prompt ordering changes between dataset versions, saved annotations can bind to the wrong prompt without any error. Prefer a stable key (session_id, explicit id, or prompt hash) and keep legacy index fallback for migration reads.

Proposed fix
+function getAnnotationKey(idx) {
+  const p = dataset[idx] || {};
+  return p.session_id || p.id || `prompt:${p.prompt || idx}`;
+}
+
 function getAnnotation(idx) {
-  return annotations[idx] || null;
+  const key = getAnnotationKey(idx);
+  return annotations[key] || annotations[idx] || null; // legacy fallback
 }
 
 function annotate(idx, tier) {
-  annotations[idx] = { tier, note: annotations[idx]?.note || '', ts: new Date().toISOString() };
+  const key = getAnnotationKey(idx);
+  annotations[key] = { tier, note: annotations[key]?.note || '', ts: new Date().toISOString() };
   saveAnnotations();
   render();
 }
 
 function saveNote(idx) {
+  const key = getAnnotationKey(idx);
   const note = document.getElementById('note-input').value;
-  if (!annotations[idx]) {
-    annotations[idx] = { tier: dataset[idx].tier || dataset[idx].llm_tier || '?', note: '', ts: new Date().toISOString() };
+  if (!annotations[key]) {
+    annotations[key] = { tier: dataset[idx].llm_tier || dataset[idx].tier || '?', note: '', ts: new Date().toISOString() };
   }
-  annotations[idx].note = note;
+  annotations[key].note = note;
   saveAnnotations();
   render();
 }
 
 function clearAnnotation(idx) {
-  delete annotations[idx];
+  const key = getAnnotationKey(idx);
+  delete annotations[key];
   saveAnnotations();
   render();
 }

Also applies to: 307-317, 323-324, 330-330, 346-347

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@router/static/visualizer.html` around lines 157 - 158, Replace the array
index-based keying mechanism for annotations with stable identifiers such as a
prompt hash, session_id, or explicit id throughout the annotation system. In the
getAnnotation function and related annotation storage/retrieval code (lines
307-324, 330, 346-347), update the key structure to use these stable identifiers
instead of array indices. Maintain backward compatibility by adding a fallback
mechanism that attempts to resolve annotations by the legacy index-based key if
the stable identifier lookup fails, allowing existing saved annotations to
continue working during migration.

Comment thread router/static/visualizer.html Outdated
Comment on lines +228 to +232
return `<div class="list-item${i === selectedIdx ? ' selected' : ''}" onclick="selectPrompt(${i})">
<div class="id">#${i} · LLM: ${displayLlm}${pred ? ' · CLS: ' + displayPred : ''}</div>
<div class="snippet">${escapeHtml(p.prompt.substring(0, 120))}</div>
<div class="tags">${tags}</div>
</div>`;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Make prompt rows keyboard-operable.

Line 228 renders selectable rows as mouse-only clickable divs. Keyboard-only users can’t focus/select prompts, which blocks the core review flow.

Proposed fix
-    return `<div class="list-item${i === selectedIdx ? ' selected' : ''}" onclick="selectPrompt(${i})">
+    return `<div
+      class="list-item${i === selectedIdx ? ' selected' : ''}"
+      role="button"
+      tabindex="0"
+      onclick="selectPrompt(${i})"
+      onkeydown="if (event.key === 'Enter' || event.key === ' ') { event.preventDefault(); selectPrompt(${i}); }"
+    >
       <div class="id">#${i} · LLM: ${displayLlm}${pred ? ' · CLS: ' + displayPred : ''}</div>
       <div class="snippet">${escapeHtml(p.prompt.substring(0, 120))}</div>
       <div class="tags">${tags}</div>
     </div>`;
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
return `<div class="list-item${i === selectedIdx ? ' selected' : ''}" onclick="selectPrompt(${i})">
<div class="id">#${i} · LLM: ${displayLlm}${pred ? ' · CLS: ' + displayPred : ''}</div>
<div class="snippet">${escapeHtml(p.prompt.substring(0, 120))}</div>
<div class="tags">${tags}</div>
</div>`;
return `<div
class="list-item${i === selectedIdx ? ' selected' : ''}"
role="button"
tabindex="0"
onclick="selectPrompt(${i})"
onkeydown="if (event.key === 'Enter' || event.key === ' ') { event.preventDefault(); selectPrompt(${i}); }"
>
<div class="id">#${i} · LLM: ${displayLlm}${pred ? ' · CLS: ' + displayPred : ''}</div>
<div class="snippet">${escapeHtml(p.prompt.substring(0, 120))}</div>
<div class="tags">${tags}</div>
</div>`;
🧰 Tools
🪛 ast-grep (0.43.0)

[warning] 211-232: Avoid assigning untrusted data to innerHTML/outerHTML or document.write
Context: listEl.innerHTML = filtered.map(i => {
const p = dataset[i];
const llm = p.llm_tier || p.tier || '?';
const pred = p.clf_tier;
const ann = getAnnotation(i);
const conflict = isConflict(i);

let tags = '';
if (conflict === true) tags += '<span class="tag conflict">CONFLICT</span>';
else if (conflict === false) tags += '<span class="tag agree">AGREE</span>';
else tags += '<span class="tag" style="background:`#333`;color:`#888`">NO DATA</span>';
if (ann) tags += '<span class="tag human">REVIEWED</span>';

const displayLlm = typeof llm === 'string' ? escapeHtml(llm.replace('agent-','')) : '?';
const displayPred = pred ? escapeHtml(pred.replace('agent-','')) : '';

return `<div class="list-item${i === selectedIdx ? ' selected' : ''}" onclick="selectPrompt(${i})">
  <div class="id">#${i} · LLM: ${displayLlm}${pred ? ' · CLS: ' + displayPred : ''}</div>
  <div class="snippet">${escapeHtml(p.prompt.substring(0, 120))}</div>
  <div class="tags">${tags}</div>
</div>`;

}).join('')
Note: [CWE-79].

(inner-outer-html)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@router/static/visualizer.html` around lines 228 - 232, The div element with
class "list-item" that renders selectable prompt rows is currently only
mouse-operable via the onclick handler for selectPrompt. To make this
keyboard-operable for accessibility, convert the div element to a button element
or add keyboard event handling and accessibility attributes (tabindex,
role="button", and onkeydown handler) to the div. This will allow keyboard-only
users to focus on and select prompts using Tab and Enter keys, enabling them to
use the core review workflow.

Comment thread scripts/benchmark_classifier.py Outdated
Comment on lines +53 to +56
for i, item in enumerate(dataset.get("prompts", [])):
prompt = item["prompt"]
expected = item["tier"]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Guard expected-label lookup to prevent benchmark crashes.

Line 55 assumes every item has tier, and Line 68 assumes it is always one of TIERS. Datasets from reclassify_all.py use llm_tier/clf_tier, and unresolved labels like "ERROR" are possible; either case can raise KeyError and stop the run.

Suggested fix
 for i, item in enumerate(dataset.get("prompts", [])):
     prompt = item["prompt"]
-    expected = item["tier"]
+    expected = item.get("llm_tier") or item.get("tier")
+    if expected not in TIERS:
+        results.append({
+            "prompt": prompt[:100],
+            "expected": expected or "UNKNOWN",
+            "predicted": "SKIPPED_UNLABELED",
+        })
+        continue
@@
-    per_tier[expected]["total"] += 1
+    per_tier[expected]["total"] += 1

Also applies to: 68-73

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/benchmark_classifier.py` around lines 53 - 56, The benchmark script
assumes every item in the prompts dataset has a "tier" key and that the value is
always one of the valid TIERS, but datasets from reclassify_all.py use
"llm_tier" or "clf_tier" instead, and may contain unresolved labels like
"ERROR". Replace the direct access to item["tier"] with a safe lookup using
.get() that checks for multiple possible key names (tier, llm_tier, clf_tier),
and add defensive validation around line 68 where the tier value is checked
against TIERS to handle unexpected values gracefully rather than raising a
KeyError that crashes the benchmark run.

Comment on lines +40 to +55
def extract_user_prompt(obs):
"""Extract and parse the raw user prompt from the observation input payload."""
inp = obs.get('input')
if not inp: return None
if isinstance(inp, str):
try: inp = json.loads(inp)
except: return None
if not isinstance(inp, dict): return None
messages = inp.get('messages', [])
if not messages: return None
for msg in reversed(messages):
if isinstance(msg, dict) and msg.get('role') == 'user':
content = msg.get('content', '')
if isinstance(content, str) and len(content.strip()) > 3:
return content.strip()
return None

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Align extract_user_prompt with the first-user-message rule.

Line 50 currently selects the last user turn, and it does not skip pseudo-user system notes. That reintroduces the prompt-selection issue already fixed in scripts/extract_prompts.py, so this script can collect the wrong text for classification.

Suggested fix
 def extract_user_prompt(obs):
     """Extract and parse the raw user prompt from the observation input payload."""
     inp = obs.get('input')
-    if not inp: return None
+    if not inp:
+        return None
     if isinstance(inp, str):
-        try: inp = json.loads(inp)
-        except: return None
-    if not isinstance(inp, dict): return None
+        try:
+            inp = json.loads(inp)
+        except json.JSONDecodeError:
+            return None
+    if not isinstance(inp, dict):
+        return None
     messages = inp.get('messages', [])
-    if not messages: return None
-    for msg in reversed(messages):
+    if not messages:
+        return None
+    for msg in messages:
         if isinstance(msg, dict) and msg.get('role') == 'user':
             content = msg.get('content', '')
             if isinstance(content, str) and len(content.strip()) > 3:
-                return content.strip()
+                stripped = content.strip()
+                if stripped.startswith('[System:') or stripped.startswith('[Note:'):
+                    continue
+                return stripped
     return None
🧰 Tools
🪛 Ruff (0.15.17)

[error] 43-43: Multiple statements on one line (colon)

(E701)


[error] 45-45: Multiple statements on one line (colon)

(E701)


[error] 46-46: Do not use bare except

(E722)


[error] 46-46: Multiple statements on one line (colon)

(E701)


[error] 47-47: Multiple statements on one line (colon)

(E701)


[error] 49-49: Multiple statements on one line (colon)

(E701)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/extract_complex.py` around lines 40 - 55, The extract_user_prompt
function currently iterates through messages in reversed order using
reversed(messages), which causes it to select the last user message instead of
the first. Change the loop to iterate forward through messages in their original
order by removing the reversed() call and updating for msg in reversed(messages)
to for msg in messages. This will ensure the function returns the first user
message, aligning with the behavior already implemented in
scripts/extract_prompts.py and preventing the prompt-selection issue mentioned
in the review.

Comment on lines +38 to +56
def extract_user_prompt(obs):
"""Extract and parse the raw user prompt from the observation input payload."""
inp = obs.get('input')
if not inp:
return None
if isinstance(inp, str):
try: inp = json.loads(inp)
except: return None
if not isinstance(inp, dict):
return None
messages = inp.get('messages', [])
if not messages:
return None
for msg in reversed(messages):
if isinstance(msg, dict) and msg.get('role') == 'user':
content = msg.get('content', '')
if isinstance(content, str) and len(content.strip()) > 3:
return content.strip()
return None

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Use the same prompt extraction semantics as the primary extractor.

Line 51 walks messages in reverse and returns the last user message, and the function does not filter pseudo-user system notes. That can reintroduce the same bad prompt capture this PR fixed elsewhere.

Suggested fix
 def extract_user_prompt(obs):
     """Extract and parse the raw user prompt from the observation input payload."""
@@
-    for msg in reversed(messages):
+    for msg in messages:
         if isinstance(msg, dict) and msg.get('role') == 'user':
             content = msg.get('content', '')
             if isinstance(content, str) and len(content.strip()) > 3:
-                return content.strip()
+                stripped = content.strip()
+                if stripped.startswith('[System:') or stripped.startswith('[Note:'):
+                    continue
+                return stripped
🧰 Tools
🪛 Ruff (0.15.17)

[error] 44-44: Multiple statements on one line (colon)

(E701)


[error] 45-45: Do not use bare except

(E722)


[error] 45-45: Multiple statements on one line (colon)

(E701)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/extract_gapfill.py` around lines 38 - 56, The extract_user_prompt
function must use the same prompt extraction semantics as the primary extractor
to avoid reintroducing the bad prompt capture issue fixed elsewhere. Currently,
the function walks through messages in reversed order looking for user role
entries without filtering pseudo-user system notes. Locate the primary
extractor's implementation in the codebase and apply its same message filtering
and validation logic to the message extraction loop in extract_user_prompt,
ensuring both functions handle user prompts consistently.

Comment thread scripts/retry_errors.py Outdated
Comment on lines +70 to +71
error_indices = [i for i, p in enumerate(dataset.get('prompts', [])) if p.get('tier') == 'ERROR']
print(f"Retrying {len(error_indices)} failed prompts (max {MAX_CHARS} chars)...")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Make retry/update logic schema-aware (tier vs clf_tier).

Line 70 only targets tier == "ERROR", while Line 95 prioritizes llm_tier in counts. With datasets written by reclassify_all.py (llm_tier + clf_tier), retries can become a no-op and post-retry counts won’t reflect classifier fixes. Also, fixed is incremented even when classify() returns "ERROR".

Suggested fix
-error_indices = [i for i, p in enumerate(dataset.get('prompts', [])) if p.get('tier') == 'ERROR']
+tier_key = "clf_tier" if any("clf_tier" in p for p in dataset.get("prompts", [])) else "tier"
+error_indices = [
+    i for i, p in enumerate(dataset.get("prompts", []))
+    if str(p.get(tier_key, "")).startswith("ERROR")
+]
@@
-            tier = classify(prompt)
+            tier = classify(prompt)
             if idx < len(prompts_list):
-                prompts_list[idx]['tier'] = tier
-            fixed += 1
+                prompts_list[idx][tier_key] = tier
+            if tier == "ERROR":
+                errors += 1
+            else:
+                fixed += 1
@@
-new_counts = Counter(p.get('llm_tier') or p.get('tier', 'ERROR') for p in dataset.get('prompts', []))
+new_counts = Counter(p.get(tier_key, 'ERROR') for p in dataset.get('prompts', []))

Also applies to: 82-98

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/retry_errors.py` around lines 70 - 71, The retry logic needs to be
schema-aware to handle both the old tier schema (single `tier` field) and the
new schema used by reclassify_all.py (which uses `llm_tier` and `clf_tier`).
Update the list comprehension at line 70 that filters for `tier == 'ERROR'` to
check for both old and new schema fields, then apply the same schema-aware logic
consistently throughout lines 82-98 when reading and counting tier values
(currently line 95 references `llm_tier` inconsistently). Additionally, the
`fixed` counter should only be incremented when the classify() function returns
a successful result (not "ERROR"), so add a check before incrementing the
counter to ensure the classification actually fixed the error.

@sheepdestroyer

Copy link
Copy Markdown
Owner Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a dual circuit breaker system, asynchronous non-blocking stats persistence, native streaming support for the agy proxy, and a dataset visualizer dashboard with corresponding backend endpoints and benchmarking scripts. The review feedback highlights several robustness improvements, including preventing potential crashes from null message content in main.py, resolving potential AttributeError and IndexError exceptions in the helper scripts, and cleaning up unused code and UI inconsistencies in the visualizer interface.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread router/main.py
# Success telemetry
latency_ms = (time.time() - start_time) * 1000.0
# Approximate prompt tokens based on messages characters
prompt_chars = sum(len(m.get("content", "")) for m in messages)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

If any message in messages has a content field explicitly set to null (which is common for assistant tool calls in OpenAI/LiteLLM formats), m.get("content", "") will return None. This will cause len(None) to raise a TypeError and crash the stream generator. Using m.get("content") or "" safely handles None values.

Suggested change
prompt_chars = sum(len(m.get("content", "")) for m in messages)
prompt_chars = sum(len(m.get("content") or "") for m in messages)

Comment on lines +128 to +135
// Load benchmark results if available
try {
const resp = await fetch('/data/benchmark_results.json');
benchmark = await resp.json();
// Map benchmark results to dataset by index
} catch(e) {
console.warn('No benchmark results yet — classifier column will show "pending"');
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The benchmark variable and the fetch call to /data/benchmark_results.json are completely unused in the visualizer. getBenchmarkPrediction reads clf_tier directly from the dataset object instead. Removing this dead code cleans up the codebase and avoids an unnecessary network request on page load.

Suggested change
// Load benchmark results if available
try {
const resp = await fetch('/data/benchmark_results.json');
benchmark = await resp.json();
// Map benchmark results to dataset by index
} catch(e) {
console.warn('No benchmark results yet — classifier column will show "pending"');
}
// Benchmark results are read directly from the dataset

let tierBtns = TIERS.map(tier => {
let cls = 'tier-btn';
if (ann && ann.tier === tier) cls += ' selected-human';
if (llm === tier && !ann) cls += ' llm-label';

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

By disabling the llm-label class when ann is present, the visual indicator for the original LLM label disappears once a human annotation is saved. However, the classifier label (classifier-label) remains visible. To maintain UI consistency and preserve context, the original LLM label should remain highlighted even after human review.

Suggested change
if (llm === tier && !ann) cls += ' llm-label';
if (llm === tier) cls += ' llm-label';

Comment thread scripts/retry_errors.py Outdated
data = json.loads(resp.read())
for m in data.get('data', []):
if 'gemma4-26a4b' in m.get('id', ''):
args = m.get('status', {}).get('args', [])

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If the status field in the model object is explicitly set to null or is not a dictionary, calling .get('args') on it will raise an AttributeError. Using a fallback dictionary or type check ensures robust parsing.

            status_obj = m.get("status") or {}\n            args = status_obj.get("args", []) if isinstance(status_obj, dict) else []

Comment thread scripts/reclassify_all.py Outdated
req = urllib.request.Request(LLAMA_SERVER_URL, data=json.dumps(payload).encode(), headers={'Content-Type':'application/json','Authorization':'Bearer local-token'})
with urllib.request.urlopen(req, timeout=30) as resp:
data = json.loads(resp.read())
return data['choices'][0]['message'].get('content','').strip()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Accessing data['choices'][0] directly without checking if choices is non-empty can raise an IndexError if the model response is empty or fails. Adding a safe fallback check prevents script crashes.

    choices = data.get("choices", [])\n    return choices[0].get("message", {}).get("content", "").strip() if choices else ""

- router/main.py: wrap lifespan yield in try/finally so stats flush always
  runs; add _background_tasks set to prevent fire-and-forget task GC (RUF006)
- router/static/visualizer.html: add keyboard accessibility (tabindex, role,
  onkeydown) to prompt list items; switch annotation keys from array index to
  stable djb2 prompt hash with backward-compatible index fallback
- scripts/benchmark_classifier.py: safe tier field lookup (tier/llm_tier/clf_tier),
  guard per_tier keying to skip ERROR/unknown labels, safe choices[0] access
- scripts/extract_complex.py: use forward iteration for first user message,
  add system-note filtering ([System: / [Note:])
- scripts/extract_gapfill.py: same as extract_complex.py
- scripts/retry_errors.py: safe status.get() to prevent AttributeError,
  schema-aware ERROR detection (tier + clf_tier), only increment fixed on success
- scripts/reclassify_all.py: safe choices[0] fallback to prevent IndexError
- test_circuit_breaker.py: replace == True/False with idiomatic assert (14 sites)
- verify_breaker.py: replace == True/False with idiomatic assert (5 sites)
@sheepdestroyer

Copy link
Copy Markdown
Owner Author

Superseded by PR #4

@sheepdestroyer
sheepdestroyer deleted the finalize-pr2-fixes-review branch June 19, 2026 12:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants